
import jieba
import pandas as pd
from neo4j import GraphDatabase


# 第二步，写匹配机制，简单：关键字匹配，先实现四种问题的查询
'''
问题类型：
1.病吃药
2.药治病
3.病什么症状
4.症状什么病
组合问题：
5.症状吃什么药

区分的方法：
区分1&2和3&4，有无药
1，2区分：病在关键词中，还是药在关键词中
3，4区分：病在关键词中，还是症状在关键词中

将问题分为两个部分，一部分是用户已知，另一部分是用户未知，从已知部分出发，通过cyber查找到未知部分
'''
def get_result(question):
    data = pd.read_csv('disease.csv')
    data = data.fillna('无')

    # question = '四肢抽搐该吃什么药'
    # '乙肝吃什么药' '乙肝有什么症状' '四肢抽搐是什么病' '定坤丹能治什么病' '四肢抽搐该吃什么药'
    for word in ['什么药','什么病','哪些药','什么症状','哪些症状','哪些表现','四肢抽搐']:
        jieba.add_word(word)
    question_key_words = set(jieba.lcut(question,cut_all = False))
    # print(question_key_words)

    # 创建匹配词典
    name_known = set(data['name'])
    alias_known = set(' '.join(list(data['alias'])).split())
    symptom_known = set((' '.join(list(data['symptom']))).split())
    drug_known = set((' '.join(list(data['drug']))).split())

    name_unknown = set(['什么病'])
    drug_unknown = set(['什么药','哪些药'])
    symptom_unknown = set(['什么症状','哪些症状','哪些表现'])


    # 病、症状、药构成三位代码
    #001 010 100
    #两个列表做交集是否为空
    # 例如001 010知道某种病，想查询症状
    name_in_table = 1 if (question_key_words & name_known) else 0
    symptom_in_table = 10 if (question_key_words & symptom_known) else 0
    drug_in_table = 100 if (question_key_words & drug_known) else 0
    known_code = name_in_table + symptom_in_table + drug_in_table

    name_in_unk = 1 if (question_key_words & name_unknown) else 0
    symptom_in_unk = 10 if (question_key_words & symptom_unknown) else 0
    drug_in_unk = 100 if (question_key_words & drug_unknown) else 0
    unknown_code = name_in_unk + symptom_in_unk + drug_in_unk
    # print(known_code,unknown_code)

    URI = "bolt://localhost:7687"
    AUTH = ("admin", "73@TuGraph")
    with GraphDatabase.driver(URI, auth=AUTH) as client:
        session = client.session(database="default")

        # 根据状态代码写查找语句
        if (known_code,unknown_code) == (1,10):

            # 知道病找症状
            name = list((question_key_words & name_known))[0]
            query = "match (n1:disease {name:\""+name+"\"})-[r:is_symptom]-(n2) return n2.name"
            ret = session.run(query)
            result = [item.get('n2.name')for item in ret.data()]

        elif (known_code,unknown_code) == (1,100):
            # 知道病想找到药
            name = list((question_key_words & name_known))[0]
            query = "match (n1:disease {name:\""+name+"\"})-[r:is_drug]-(n2) return n2.name"
            ret = session.run(query)
            result = [item.get('n2.name')for item in ret.data()]

        elif (known_code,unknown_code) == (10,1):
            # 知道症状找病
            symptom = list((question_key_words & symptom_known))[0]
            query = "match (n1:disease)-[r:is_symptom]-(n2:symptom{name:\""+symptom+"\"}) return n1.name"
            ret = session.run(query)
            result = [item.get('n1.name')for item in ret.data()]


        elif (known_code,unknown_code) == (100,1):
            # 知道药找病
            drug = list((question_key_words & drug_known))[0]
            query = "match (n1:disease)-[r:is_drug]-(n2:drug{name:\""+drug+"\"}) return n1.name"
            ret = session.run(query)
            result = [item.get('n1.name')for item in ret.data()]

        elif (known_code,unknown_code) == (10,100):
            # 知道症状找药
            symptom = list((question_key_words & symptom_known))[0]
            query_1 = "match (n1:disease)-[r:is_symptom]-(n2:symptom{name:\""+symptom+"\"}) return n1.name"
            ret = session.run(query_1)
            diseases = [item.get('n1.name') for item in ret.data()]

            result = pd.DataFrame(columns=['name', 'drug'])
            for name in diseases:
                query = "match (n1:disease {name:\""+name+"\"})-[r:is_drug]-(n2) return n2.name"
                ret = session.run(query)
                drugs = [item.get('n2.name') for item in ret.data()]
                rs = pd.DataFrame({'name': name ,'drug': drugs})
                result = pd.concat([result,rs], axis=0, ignore_index=True)

        # print(result)
        return result